package com.kuaiyin.util;

import com.kuaiyin.enums.BizCodeEnum;
import com.kuaiyin.exception.BaseException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.FileUtils;
import org.apache.jmeter.engine.StandardJMeterEngine;
import org.apache.jmeter.reporters.ResultCollector;
import org.apache.jmeter.save.SaveService;
import org.apache.jmeter.testelement.TestElement;
import org.apache.jorphan.collections.HashTree;
import org.apache.jorphan.collections.ListedHashTree;

import java.io.File;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.*;

@Slf4j
public class JMeterExecutorUtil {

    // 固定线程池
    private static final ExecutorService executor = Executors.newFixedThreadPool(5);
    // 缓存正在运行的 JMeter 引擎实例
    private static final Map<String, StandardJMeterEngine> runningEngines = new ConcurrentHashMap<>();

    /**
     * 异步执行JMX，返回Future，可用于等待结果
     */
    public static Future<Map<String, Object>> executeJmxAsync(String executionId,String jmxContent) {
        return executor.submit(() -> {
            try {
                return executeJmxAndWait(executionId,jmxContent);
            } catch (Exception e) {
                log.error("异步执行JMeter出现异常", e);
                throw new RuntimeException(e);
            }
        });
    }



    /**
     * 执行 JMeter 测试：将传入的 JMX 内容写入临时文件并运行，等待执行结束后解析结果。
     * @param jmxContent JMeter 测试计划（JMX）的 XML 字符串
     * @throws IOException 如果文件操作失败
     * @throws InterruptedException 如果线程等待被中断
     */
    private static Map<String, Object> executeJmxAndWait(String executionId,String jmxContent) throws IOException, InterruptedException {
        File tempFile = null;
        File tempJtlFile = null;
        StandardJMeterEngine engine = null;
        // 用于等待 JMeter 测试执行完毕
        CountDownLatch latch = new CountDownLatch(1);


        try {
            if (!isJMeterJmxContent(jmxContent)) {
                throw BaseException.of(BizCodeEnum.INVALID_PARAMETER);
            }
            // 1. 创建临时 JMX 文件，并写入传入的 jmxContent 测试计划
            tempFile = File.createTempFile("jmeter_", ".jmx");
            FileUtils.writeStringToFile(tempFile, jmxContent, "UTF-8");

            // 2. 创建临时 JTL 文件（CSV 格式），用于保存 JMeter 测试执行结果
            tempJtlFile = File.createTempFile("jmeter_result_", ".jtl");

            // 3. 初始化 JMeter 环境（加载 jmeter.properties 等配置）
            JmeterUtil.initJmeterProperties();

            // 4. 加载测试计划到 HashTree 结构
            HashTree testPlanTree = SaveService.loadTree(tempFile);
            ListedHashTree listedHashTree = new ListedHashTree();
            // 深拷贝到新的 HashTree，防止后续修改影响原始结构
            deepCopy(testPlanTree, listedHashTree);

            // 5. 创建 Summariser（控制台汇总摘要）和 ResultCollector（结果收集器）
            ResultCollector summaryLogger = new ResultCollector();
            summaryLogger.setName("summary");
            summaryLogger.setFilename(tempJtlFile.getAbsolutePath());

            // 6. 将结果收集器添加到 ThreadGroup 节点下
            for (Object testPlanObj : listedHashTree.keySet()) {
                HashTree planTree = listedHashTree.getTree(testPlanObj);
                for (Object childObj : planTree.keySet()) {
                    if (childObj instanceof TestElement element
                            && "ThreadGroup".equals(element.getClass().getSimpleName())) {
                        HashTree threadGroupTree = planTree.getTree(childObj);
                        threadGroupTree.add(summaryLogger);

                    }
                }
            }

            // 7. 初始化并配置 JMeter 引擎
            engine = new StandardJMeterEngine();
            // 注册引擎实例
            runningEngines.put(executionId, engine);
            engine.configure(listedHashTree);

            // 8. 创建一个线程监控测试执行状态：检测引擎是否还在运行
            StandardJMeterEngine finalEngine = engine;
            Thread monitor = new Thread(() -> {
                try {
                    while (finalEngine.isActive()) {
                        Thread.sleep(500); // 每 0.5 秒检查一次
                    }
                    latch.countDown(); // 测试结束，通知主线程继续
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            });
            monitor.start();

            // 9. 启动 JMeter 测试
            engine.run();

            // 10. 等待测试执行结束
            latch.await();
            // 11. 测试结束后，解析生成的 JTL 文件，获取汇总统计信息
            return JtlParserUtil.parseJtl(tempJtlFile);

        } catch (Exception e) {
            // 12. 捕获异常并记录日志，同时抛出运行时异常
            log.error("JMeter执行异常", e);
            throw new RuntimeException("JMeter执行异常", e);
        } finally {
            if (engine != null) {
                try {
                    engine.stopTest(true); // 防御性停止
                    engine.reset();        // 清理资源
                } catch (Exception ex) {
                    log.warn("关闭 JMeter 引擎时异常", ex);
                }
            }

            // 13. 删除临时文件，避免磁盘堆积
            if (tempFile != null && tempFile.exists()) {
                boolean deleted = tempFile.delete();
                if (!deleted) tempFile.deleteOnExit();
            }
            if (tempJtlFile != null && tempJtlFile.exists()) tempJtlFile.deleteOnExit();
        }
    }

    /**
     * 停止指定的 JMeter 测试
     * @param executionId 执行 ID
     */
    public static void stopJmxExecution(String executionId) {
        StandardJMeterEngine engine = runningEngines.get(executionId);
        if (engine != null) {
            try {
                log.info("正在停止 JMeter 测试: {}", executionId);
                engine.stopTest(true);
                engine.reset();
            } finally {
                runningEngines.remove(executionId);
            }
        } else {
            log.warn("未找到对应的 JMeter 引擎实例: {}", executionId);
        }
    }

    /**
     * HashTree 深拷贝
     */
    private static void deepCopy(HashTree source, ListedHashTree target) {
        for (Object key : source.keySet()) {
            target.add(key);
            HashTree child = source.getTree(key);
            if (!child.isEmpty()) {
                deepCopy(child, target.getTree(key));
            }
        }
    }

    public static boolean isJMeterJmxContent(String xmlContent) {
        if (xmlContent == null || xmlContent.trim().isEmpty()) {
            return false;
        }
        String trimmed = xmlContent.trim();
        // 检查是否有 <?xml 开头 和 <TestPlan> 标签
        return trimmed.startsWith("<?xml") && trimmed.contains("<TestPlan");
    }

}
